package cn.iocoder.yudao.framework.support;

import com.baomidou.dynamic.datasource.annotation.DS;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.aop.framework.AopProxyUtils;
import org.springframework.core.BridgeMethodResolver;
import org.springframework.core.MethodClassKey;
import org.springframework.core.annotation.AnnotatedElementUtils;
import org.springframework.core.annotation.AnnotationAttributes;
import org.springframework.util.ClassUtils;

import java.lang.reflect.AnnotatedElement;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.lang.reflect.Proxy;
import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * TODO
 *
 * @className: DataSourceClassResolver
 * @author: tangkh
 * @date: 2023-09-19 18:02
 */
public class DataSourceClassResolver {
    private static final Logger log = LoggerFactory.getLogger(DataSourceClassResolver.class);
    private static boolean mpEnabled = false;
    private static Field mapperInterfaceField;
    private final Map<Object, String> dsCache = new ConcurrentHashMap<>();
    private final boolean allowedPublicOnly;

    public DataSourceClassResolver(boolean allowedPublicOnly) {
        this.allowedPublicOnly = allowedPublicOnly;
    }

    public String findKey(Method method, Object targetObject) {
        if (method.getDeclaringClass() == Object.class) {
            return "";
        } else {
            Object cacheKey = new MethodClassKey(method, targetObject.getClass());
            String ds = this.dsCache.get(cacheKey);
            if (ds == null) {
                ds = this.computeDatasource(method, targetObject);
                if (ds == null) {
                    ds = "";
                }

                this.dsCache.put(cacheKey, ds);
            }

            return ds;
        }
    }

    private String computeDatasource(Method method, Object targetObject) {
        if (this.allowedPublicOnly && !Modifier.isPublic(method.getModifiers())) {
            return null;
        } else {
            String dsAttr = this.findDataSourceAttribute(method);
            if (dsAttr != null) {
                return dsAttr;
            } else {
                Class<?> targetClass = targetObject.getClass();
                Class<?> userClass = ClassUtils.getUserClass(targetClass);
                Method specificMethod = ClassUtils.getMostSpecificMethod(method, userClass);
                specificMethod = BridgeMethodResolver.findBridgedMethod(specificMethod);
                dsAttr = this.findDataSourceAttribute(specificMethod);
                if (dsAttr != null) {
                    return dsAttr;
                } else {
                    dsAttr = this.findDataSourceAttribute(userClass);
                    if (dsAttr != null && ClassUtils.isUserLevelMethod(method)) {
                        return dsAttr;
                    } else {
                        Iterator var7 = ClassUtils.getAllInterfacesForClassAsSet(userClass).iterator();

                        do {
                            if (!var7.hasNext()) {
                                if (specificMethod != method) {
                                    dsAttr = this.findDataSourceAttribute(method);
                                    if (dsAttr != null) {
                                        return dsAttr;
                                    }

                                    dsAttr = this.findDataSourceAttribute(method.getDeclaringClass());
                                    if (dsAttr != null && ClassUtils.isUserLevelMethod(method)) {
                                        return dsAttr;
                                    }
                                }

                                return this.getDefaultDataSourceAttr(targetObject);
                            }

                            Class<?> interfaceClazz = (Class)var7.next();
                            dsAttr = this.findDataSourceAttribute(interfaceClazz);
                        } while(dsAttr == null);

                        return dsAttr;
                    }
                }
            }
        }
    }

    private String getDefaultDataSourceAttr(Object targetObject) {
        Class<?> targetClass = targetObject.getClass();
        Class clazz;
        String datasourceAttr;
        if (!Proxy.isProxyClass(targetClass)) {
            for(clazz = targetClass; clazz != Object.class; clazz = clazz.getSuperclass()) {
                datasourceAttr = this.findDataSourceAttribute(clazz);
                if (datasourceAttr != null) {
                    return datasourceAttr;
                }
            }
        }

        if (mpEnabled) {
            clazz = this.getMapperInterfaceClass(targetObject);
            if (clazz != null) {
                datasourceAttr = this.findDataSourceAttribute(clazz);
                if (datasourceAttr != null) {
                    return datasourceAttr;
                }

                return this.findDataSourceAttribute(clazz.getSuperclass());
            }
        }

        return null;
    }

    private Class<?> getMapperInterfaceClass(Object target) {
        Object current;
        Object currentRefObject;
        for(current = target; Proxy.isProxyClass(current.getClass()); current = currentRefObject) {
            currentRefObject = AopProxyUtils.getSingletonTarget(current);
            if (currentRefObject == null) {
                break;
            }
        }

        try {
            if (Proxy.isProxyClass(current.getClass())) {
                return (Class)mapperInterfaceField.get(Proxy.getInvocationHandler(current));
            }
        } catch (IllegalAccessException var4) {
        }

        return null;
    }

    private String findDataSourceAttribute(AnnotatedElement ae) {
        if (ae == null) {
            return null;
        } else {
            AnnotationAttributes attributes = AnnotatedElementUtils.getMergedAnnotationAttributes(ae, DS.class);
            return attributes != null ? attributes.getString("value") : null;
        }
    }

    static {
        Class proxyClass = null;

        try {
            proxyClass = Class.forName("com.baomidou.mybatisplus.core.override.MybatisMapperProxy");
        } catch (ClassNotFoundException var7) {
            try {
                proxyClass = Class.forName("com.baomidou.mybatisplus.core.override.PageMapperProxy");
            } catch (ClassNotFoundException var6) {
                try {
                    proxyClass = Class.forName("org.apache.ibatis.binding.MapperProxy");
                } catch (ClassNotFoundException var5) {
                }
            }
        }

        if (proxyClass != null) {
            try {
                mapperInterfaceField = proxyClass.getDeclaredField("mapperInterface");
                mapperInterfaceField.setAccessible(true);
                mpEnabled = true;
            } catch (NoSuchFieldException var4) {
                var4.printStackTrace();
            }
        }
    }
}
